#!/usr/bin/env python
# -*- coding: utf-8 -*-

from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton, 
                            QListWidget, QListWidgetItem, QLabel, QComboBox,
                            QToolBar, QAction, QFileDialog, QSplitter, 
                            QScrollArea, QFrame, QColorDialog, QInputDialog,
                            QMessageBox, QGroupBox, QFormLayout, QLineEdit,
                            QSpinBox, QDoubleSpinBox, QCheckBox, QDialog, QProgressBar,
                            QProgressDialog, QScrollBar, QApplication, QMenu)
from PyQt5.QtCore import Qt, QSize, QRect, QPoint, pyqtSignal, QTimer
from PyQt5.QtGui import QIcon, QPixmap, QPainter, QPen, QColor, QImage, QCursor, QFont
import os  # 添加os模块导入
import glob  # 添加glob模块导入

class AnnotationCanvas(QLabel):
    """标注画布组件"""
    
    # 自定义信号
    box_created = pyqtSignal(QRect, str)  # 创建了新的标注框
    
    def __init__(self):
        super().__init__()
        
        # 设置画布属性
        self.setAlignment(Qt.AlignCenter)
        self.setMinimumSize(600, 400)
        self.setFrameStyle(QFrame.StyledPanel)
        self.setMouseTracking(True)  # 启用鼠标跟踪
        
        # 启用接收键盘事件的能力
        self.setFocusPolicy(Qt.StrongFocus)
        
        # 初始化变量
        self.image = None
        self.pixmap = None
        self.drawing = False
        self.current_rect = QRect()
        self.current_class = "person"  # 默认类别
        self.boxes = []  # 已标注的框列表，每个元素为(rect, class_name, color)元组
        self.colors = {}  # 类别对应的颜色
        
        # 缩放相关设置
        self.scale_factor = 1.0  # 缩放因子
        self.min_scale_factor = 0.1  # 最小缩放比例
        self.max_scale_factor = 10.0  # 最大缩放比例
        self.zoom_step = 0.1  # 每次缩放的步长
        self.original_pixmap_size = QSize()
        
        # 拖动相关设置
        self.panning = False
        self.pan_start_pos = QPoint()
        self.last_pan_pos = QPoint()
        self.setCursor(Qt.ArrowCursor)
        
        # 直接拖动偏移
        self.offset = QPoint(0, 0)  # 图像偏移量
        
        # 设置默认颜色
        self.colors["person"] = QColor(255, 0, 0)  # 红色
        self.colors["car"] = QColor(0, 255, 0)  # 绿色
        self.colors["dog"] = QColor(0, 0, 255)  # 蓝色
        
    def set_image(self, image_path):
        """设置要标注的图像"""
        self.image_path = image_path
        self.image = QImage(image_path)
        if self.image.isNull():
            QMessageBox.warning(self, "错误", f"无法加载图像: {image_path}")
            return False
        
        # 清除之前的标注
        self.boxes = []
        
        self.pixmap = QPixmap.fromImage(self.image)
        self.original_pixmap_size = self.pixmap.size()
        self.scale_factor = 1.0  # 重置缩放因子
        self.update_pixmap()
        
        return True
    
    def update_pixmap(self):
        """更新显示的图像，包括所有标注"""
        if self.pixmap is None:
            return
        
        # 创建工作副本
        working_pixmap = QPixmap(self.pixmap)
        painter = QPainter(working_pixmap)
        
        # 绘制所有已有的框
        for rect, class_name, color in self.boxes:
            pen = QPen(color, 2)
            painter.setPen(pen)
            painter.drawRect(rect)
            
            # 绘制类别标签
            text_rect = QRect(rect.x(), rect.y() - 20, rect.width(), 20)
            painter.fillRect(text_rect, QColor(0, 0, 0, 128))  # 半透明背景
            painter.setPen(Qt.white)
            painter.drawText(text_rect, Qt.AlignCenter, class_name)
            
        # 如果正在绘制，也绘制当前的框
        if self.drawing and not self.current_rect.isNull():
            pen = QPen(self.colors.get(self.current_class, Qt.red), 2)
            pen.setStyle(Qt.DashLine)  # 绘制虚线框
            painter.setPen(pen)
            painter.drawRect(self.current_rect)
            
        painter.end()
        
        # 保存旧的尺寸用于比较
        old_size = getattr(self, 'displayed_pixmap_size', QSize()) if hasattr(self, 'displayed_pixmap_size') else QSize()
        
        # 应用缩放
        scaled_pixmap = working_pixmap.scaled(
            int(self.original_pixmap_size.width() * self.scale_factor),
            int(self.original_pixmap_size.height() * self.scale_factor),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        
        # 创建叠加层显示信息
        overlay_painter = QPainter(scaled_pixmap)
        overlay_painter.setRenderHint(QPainter.Antialiasing)
        overlay_painter.setRenderHint(QPainter.TextAntialiasing)
        
        # 显示所有需要的信息
        info_texts = []
        
        # 添加缩放信息
        if self.scale_factor != 1.0:
            info_texts.append(f"缩放: {int(self.scale_factor * 100)}%")
        
        # 添加拖动模式信息
        if self.panning:
            info_texts.append("拖动模式 (拖动中)")
            
        # 添加偏移量信息（用于调试）
        if self.offset != QPoint(0, 0):
            info_texts.append(f"偏移量: ({self.offset.x()}, {self.offset.y()})")
            
        # 添加使用提示（只在第一次渲染时显示）
        if not hasattr(self, 'showed_tips'):
            info_texts.append("提示: 鼠标右键可拖动图像")
            info_texts.append("提示: 滚轮可缩放图像")
            self.showed_tips = True
            
            # 创建定时器，5秒后清除提示
            QTimer.singleShot(5000, lambda: setattr(self, 'showed_tips', False) or self.update_pixmap())
            
        # 显示所有信息文本
        y_offset = 10
        for text in info_texts:
            # 计算文本框大小
            fm = overlay_painter.fontMetrics()
            text_width = fm.width(text) + 20
            text_height = 30
            
            # 绘制背景和文本
            info_rect = QRect(10, y_offset, text_width, text_height)
            overlay_painter.fillRect(info_rect, QColor(0, 0, 0, 160))
            overlay_painter.setPen(Qt.white)
            overlay_painter.setFont(QFont("Arial", 10, QFont.Bold))
            overlay_painter.drawText(info_rect, Qt.AlignCenter, text)
            
            # 更新下一个文本的位置
            y_offset += text_height + 5
            
        overlay_painter.end()
        
        # 设置画布的像素图
        self.setPixmap(scaled_pixmap)
        # 记录显示的pixmap尺寸
        self.displayed_pixmap_size = scaled_pixmap.size()
        
        # 重绘以应用偏移
        self.update()
        
    def set_current_class(self, class_name):
        """设置当前正在标注的类别"""
        self.current_class = class_name
    
    def zoom_in(self, center_pos=None):
        """放大图像
        
        参数:
        center_pos - 缩放中心点，默认为None表示图像中心
        """
        if self.scale_factor >= self.max_scale_factor:
            return
            
        # 获取新的缩放因子
        old_scale = self.scale_factor
        new_scale = old_scale + self.zoom_step
        new_scale = min(new_scale, self.max_scale_factor)
        
        # 如果指定了中心点，需要保持中心点不变
        if center_pos is not None:
            # 保存鼠标位置对应的图像点
            image_point = self.get_image_coordinates(center_pos)
            
            # 更新缩放因子
            self.scale_factor = new_scale
            
            # 获取画布的大小
            label_rect = self.rect()
            
            # 缩放后图像尺寸
            scaled_width = int(self.original_pixmap_size.width() * new_scale)
            scaled_height = int(self.original_pixmap_size.height() * new_scale)
            
            # 计算新的偏移量，使鼠标下的图像点保持不变
            new_center_x = center_pos.x() - image_point.x() * new_scale
            new_center_y = center_pos.y() - image_point.y() * new_scale
            
            # 考虑画布中心位置
            canvas_center_x = (label_rect.width() - scaled_width) / 2
            canvas_center_y = (label_rect.height() - scaled_height) / 2
            
            # 设置新的偏移量
            self.offset.setX(int(new_center_x - canvas_center_x))
            self.offset.setY(int(new_center_y - canvas_center_y))
        else:
            # 不指定中心点，直接更新缩放因子
            self.scale_factor = new_scale
        
        # 更新显示
        self.update_pixmap()
        
    def zoom_out(self, center_pos=None):
        """缩小图像
        
        参数:
        center_pos - 缩放中心点，默认为None表示图像中心
        """
        if self.scale_factor <= self.min_scale_factor:
            return
            
        # 获取新的缩放因子
        old_scale = self.scale_factor
        new_scale = old_scale - self.zoom_step
        new_scale = max(new_scale, self.min_scale_factor)
        
        # 如果指定了中心点，需要保持中心点不变
        if center_pos is not None:
            # 保存鼠标位置对应的图像点
            image_point = self.get_image_coordinates(center_pos)
            
            # 更新缩放因子
            self.scale_factor = new_scale
            
            # 获取画布的大小
            label_rect = self.rect()
            
            # 缩放后图像尺寸
            scaled_width = int(self.original_pixmap_size.width() * new_scale)
            scaled_height = int(self.original_pixmap_size.height() * new_scale)
            
            # 计算新的偏移量，使鼠标下的图像点保持不变
            new_center_x = center_pos.x() - image_point.x() * new_scale
            new_center_y = center_pos.y() - image_point.y() * new_scale
            
            # 考虑画布中心位置
            canvas_center_x = (label_rect.width() - scaled_width) / 2
            canvas_center_y = (label_rect.height() - scaled_height) / 2
            
            # 设置新的偏移量
            self.offset.setX(int(new_center_x - canvas_center_x))
            self.offset.setY(int(new_center_y - canvas_center_y))
        else:
            # 不指定中心点，直接更新缩放因子
            self.scale_factor = new_scale
        
        # 更新显示
        self.update_pixmap()
        
    def reset_zoom(self):
        """重置缩放"""
        self.scale_factor = 1.0
        self.offset = QPoint(0, 0)  # 重置偏移量
        self.update_pixmap()

    def wheelEvent(self, event):
        """处理鼠标滚轮事件，实现缩放"""
        if self.pixmap is None:
            return
            
        # 获取滚轮增量
        delta = event.angleDelta().y()
        
        # 获取当前鼠标位置在窗口中的坐标
        mouse_pos = event.pos()
        
        # 计算新的缩放比例
        if delta > 0:  # 放大
            new_scale = min(self.scale_factor + self.zoom_step, self.max_scale_factor)
        else:  # 缩小
            new_scale = max(self.scale_factor - self.zoom_step, self.min_scale_factor)
            
        # 如果缩放比例未改变，则不处理
        if new_scale == self.scale_factor:
            event.accept()
            return
            
        # 保存鼠标位置对应的图像点
        image_point = self.get_image_coordinates(mouse_pos)
        
        # 更新缩放因子
        self.scale_factor = new_scale
        
        # 获取画布的大小
        label_rect = self.rect()
        
        # 缩放后图像尺寸
        scaled_width = int(self.original_pixmap_size.width() * new_scale)
        scaled_height = int(self.original_pixmap_size.height() * new_scale)
        
        # 计算新的偏移量，使鼠标下的图像点保持不变
        new_center_x = mouse_pos.x() - image_point.x() * new_scale
        new_center_y = mouse_pos.y() - image_point.y() * new_scale
        
        # 考虑画布中心位置
        canvas_center_x = (label_rect.width() - scaled_width) / 2
        canvas_center_y = (label_rect.height() - scaled_height) / 2
        
        # 设置新的偏移量
        self.offset.setX(int(new_center_x - canvas_center_x))
        self.offset.setY(int(new_center_y - canvas_center_y))
        
        # 更新显示
        self.update_pixmap()
        
        # 处理事件，不传递给父窗口
        event.accept()

    def mousePressEvent(self, event):
        """鼠标按下事件处理"""
        if event.button() == Qt.LeftButton and self.pixmap is not None:
            self.drawing = True
            
            # 考虑缩放后的坐标转换
            pos = self.get_image_coordinates(event.pos())
            self.start_point = pos
            self.current_rect = QRect(self.start_point, QSize(0, 0))
        elif (event.button() == Qt.MiddleButton or event.button() == Qt.RightButton) and self.pixmap is not None:
            # 中键或右键拖动图片
            self.panning = True
            self.pan_start_pos = event.pos()
            self.setCursor(Qt.ClosedHandCursor)  # 更改光标为抓手
            
            # 更新显示状态
            self.update_pixmap()
            
            # 接受事件
            event.accept()
            
    def mouseMoveEvent(self, event):
        """鼠标移动事件处理"""
        if self.drawing and self.pixmap is not None:
            # 考虑缩放后的坐标转换
            pos = self.get_image_coordinates(event.pos())
            
            # 直接使用起点和当前点创建矩形，让QRect自己处理归一化
            width = pos.x() - self.start_point.x()
            height = pos.y() - self.start_point.y()
            self.current_rect = QRect(self.start_point.x(), self.start_point.y(), width, height)
            
            # 重绘
            self.update_pixmap()
        elif self.panning and self.pixmap is not None:
            # 获取当前鼠标位置
            current_pos = event.pos()
            
            # 计算移动距离
            dx = current_pos.x() - self.pan_start_pos.x()
            dy = current_pos.y() - self.pan_start_pos.y()
            
            # 直接更新偏移量
            self.offset.setX(self.offset.x() + dx)
            self.offset.setY(self.offset.y() + dy)
            
            # 更新拖动起始点
            self.pan_start_pos = current_pos
            
            # 重绘
            self.update()
            
            # 接受事件
            event.accept()
            
    def mouseReleaseEvent(self, event):
        """鼠标释放事件处理"""
        if event.button() == Qt.LeftButton and self.drawing:
            self.drawing = False
            
            # 考虑缩放后的坐标转换
            pos = self.get_image_coordinates(event.pos())
            
            # 直接使用起点和当前点创建矩形
            width = pos.x() - self.start_point.x()
            height = pos.y() - self.start_point.y()
            self.current_rect = QRect(self.start_point.x(), self.start_point.y(), width, height)
            
            # 如果矩形是反向的（宽度或高度为负），则进行标准化
            self.current_rect = self.current_rect.normalized()
            
            # 确保矩形有一定大小
            if self.current_rect.width() > 5 and self.current_rect.height() > 5:
                # 添加到框列表
                color = self.colors.get(self.current_class, Qt.red)
                self.boxes.append((self.current_rect, self.current_class, color))
                
                # 发出信号
                self.box_created.emit(self.current_rect, self.current_class)
                
                # 重绘
                self.update_pixmap()
            
            # 重置当前矩形
            self.current_rect = QRect()
        elif (event.button() == Qt.MiddleButton or event.button() == Qt.RightButton) and self.panning:
            # 结束图片拖动
            self.panning = False
            self.setCursor(Qt.ArrowCursor)  # 恢复默认光标
            
            # 更新显示状态
            self.update_pixmap()
            
            # 接受事件
            event.accept()

    def get_scroll_area(self):
        """获取父级ScrollArea"""
        parent = self.parent()
        while parent is not None:
            if isinstance(parent, QScrollArea):
                return parent
            parent = parent.parent()
        return None

    def get_image_coordinates(self, pos):
        """将窗口坐标转换为图像坐标，考虑缩放因子和位置偏移"""
        if self.pixmap is None:
            return pos
        
        # 获取Label的实际大小
        label_rect = self.rect()
        
        # 计算缩放后的实际图像尺寸
        scaled_width = int(self.original_pixmap_size.width() * self.scale_factor)
        scaled_height = int(self.original_pixmap_size.height() * self.scale_factor)
        
        # 计算图像在视图中的中心位置
        center_x = (label_rect.width() - scaled_width) / 2 + self.offset.x()
        center_y = (label_rect.height() - scaled_height) / 2 + self.offset.y()
        
        # 将窗口坐标转换为相对于图像左上角的坐标
        image_x = (pos.x() - center_x) / self.scale_factor
        image_y = (pos.y() - center_y) / self.scale_factor
        
        # 确保坐标在有效范围内
        image_x = max(0, min(image_x, self.original_pixmap_size.width() - 1))
        image_y = max(0, min(image_y, self.original_pixmap_size.height() - 1))
        
        return QPoint(int(image_x), int(image_y))

    def clear_annotations(self):
        """清除所有标注"""
        self.boxes = []
        self.update_pixmap()
        
    def remove_last_annotation(self):
        """移除最后一个标注"""
        if self.boxes:
            self.boxes.pop()
            self.update_pixmap()
            
    def export_annotations(self):
        """导出标注为YOLO格式"""
        if not self.boxes or self.image is None:
            return []
            
        results = []
        img_width = self.image.width()
        img_height = self.image.height()
        
        # 获取类别索引映射
        class_indices = {}
        parent = self.parent()
        while parent:
            if hasattr(parent, 'class_combo'):
                # 找到了包含class_combo的父组件
                for i in range(parent.class_combo.count()):
                    class_name = parent.class_combo.itemText(i)
                    class_indices[class_name] = i
                break
            parent = parent.parent()
        
        # 如果没有找到父级的class_combo，使用默认映射
        if not class_indices:
            class_indices = {"person": 0, "car": 1, "dog": 2}
        
        for rect, class_name, _ in self.boxes:
            # 将矩形坐标转换为YOLO格式（归一化中心点坐标和宽高）
            x_center = (rect.x() + rect.width() / 2) / img_width
            y_center = (rect.y() + rect.height() / 2) / img_height
            width = rect.width() / img_width
            height = rect.height() / img_height
            
            # 使用动态类别索引
            class_idx = class_indices.get(class_name, 0)
            
            # YOLO格式：<class_idx> <x_center> <y_center> <width> <height>
            results.append(f"{class_idx} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
            
        return results

    def resizeEvent(self, event):
        """窗口大小变化事件"""
        super().resizeEvent(event)
        self.update_pixmap()
        
    def paintEvent(self, event):
        """重写绘制事件"""
        # 使用QLabel的pixmap方法获取当前显示的pixmap
        current_pixmap = super().pixmap()
        if current_pixmap is None:
            super().paintEvent(event)
            return
            
        # 创建画家
        painter = QPainter(self)
        
        # 计算居中显示的位置
        x = (self.width() - current_pixmap.width()) / 2 + self.offset.x()
        y = (self.height() - current_pixmap.height()) / 2 + self.offset.y()
        
        # 绘制图像
        painter.drawPixmap(int(x), int(y), current_pixmap)
        painter.end()

    def keyPressEvent(self, event):
        """处理键盘事件并转发给父组件"""
        # 如果是特定的键盘操作，先尝试由父组件处理
        parent = self.parent()
        while parent:
            if hasattr(parent, 'keyPressEvent'):
                # 找到了AnnotationPanel父组件
                parent.keyPressEvent(event)
                if event.isAccepted():
                    return
                break
            parent = parent.parent()
        
        # 如果父组件没有处理这个事件，则调用基类方法
        super().keyPressEvent(event)
        
    def focusInEvent(self, event):
        """当控件获得焦点时"""
        super().focusInEvent(event)
        # 可以在这里添加视觉提示，表明画布当前有焦点
        
    def focusOutEvent(self, event):
        """当控件失去焦点时"""
        super().focusOutEvent(event)
        # 可以在这里清除焦点视觉提示


class AnnotationPanel(QWidget):
    """数据标注面板"""
    
    def __init__(self):
        super().__init__()
        self.init_ui()
        # 设置图像索引和文件列表
        self.current_index = -1
        self.image_files = []
        self.current_dir = ""
        # 设置焦点策略，确保组件能接收键盘事件
        self.setFocusPolicy(Qt.StrongFocus)
        
    def init_ui(self):
        """初始化UI组件"""
        main_layout = QVBoxLayout(self)
        
        # 工具栏
        toolbar = QToolBar()
        toolbar.setIconSize(QSize(24, 24))
        toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
        # 设置无边框样式
        toolbar.setStyleSheet("""
            QToolBar { 
                border: none; 
                background: transparent; 
                spacing: 8px; 
            }
            QToolButton {
                border: none;
                border-radius: 4px;
                padding: 4px;
                background: transparent;
            }
            QToolButton:hover {
                background-color: rgba(200, 200, 200, 50);
            }
            QToolButton:pressed {
                background-color: rgba(180, 180, 180, 100);
            }
        """)
        
        # 添加"加载项目图片"按钮
        load_project_images_action = QAction("加载项目图片", self)
        load_project_images_action.triggered.connect(self.load_project_images)
        toolbar.addAction(load_project_images_action)
        
        # 打开图像按钮 (Ctrl+O)
        open_action = QAction("打开图像 (Ctrl+O)", self)
        open_action.setShortcut("Ctrl+O")
        open_action.triggered.connect(self.open_image)
        toolbar.addAction(open_action)
        
        # 批量标注按钮 (Ctrl+B)
        batch_action = QAction("批量标注 (Ctrl+B)", self)
        batch_action.setShortcut("Ctrl+B")
        batch_action.triggered.connect(self.batch_annotate)
        toolbar.addAction(batch_action)
        
        # 保存标注按钮 (Ctrl+S)
        save_action = QAction("保存标注 (Ctrl+S)", self)
        save_action.setShortcut("Ctrl+S")
        save_action.triggered.connect(self.save_annotations)
        toolbar.addAction(save_action)
        
        # 清除标注按钮 (Ctrl+Del)
        clear_action = QAction("清除标注 (Ctrl+Del)", self)
        clear_action.setShortcut("Ctrl+Del")
        clear_action.triggered.connect(self.clear_annotations)
        toolbar.addAction(clear_action)
        
        # 撤销按钮 (Ctrl+Z)
        undo_action = QAction("撤销 (Ctrl+Z)", self)
        undo_action.setShortcut("Ctrl+Z")
        undo_action.triggered.connect(self.undo_annotation)
        toolbar.addAction(undo_action)
        
        # 添加分隔符
        toolbar.addSeparator()
        
        # 添加图像缩放控制按钮
        zoom_in_action = QAction("放大 (+)", self)
        zoom_in_action.setShortcut("+")
        zoom_in_action.triggered.connect(self.zoom_in_image)
        toolbar.addAction(zoom_in_action)
        
        zoom_out_action = QAction("缩小 (-)", self)
        zoom_out_action.setShortcut("-")
        zoom_out_action.triggered.connect(self.zoom_out_image)
        toolbar.addAction(zoom_out_action)
        
        reset_zoom_action = QAction("重置缩放 (0)", self)
        reset_zoom_action.setShortcut("0")
        reset_zoom_action.triggered.connect(self.reset_zoom)
        toolbar.addAction(reset_zoom_action)
        
        # 添加分隔符
        toolbar.addSeparator()
        
        # 上一张/下一张按钮 (A/D)
        prev_action = QAction("上一张 (A)", self)
        prev_action.setShortcut("A")
        prev_action.triggered.connect(self.load_prev_image)
        toolbar.addAction(prev_action)
        
        next_action = QAction("下一张 (D)", self)
        next_action.setShortcut("D")
        next_action.triggered.connect(self.load_next_image)
        toolbar.addAction(next_action)
        
        # 保存按钮引用
        self.prev_action = prev_action
        self.next_action = next_action
        
        # 添加工具栏
        main_layout.addWidget(toolbar)
        
        # 创建拆分器
        splitter = QSplitter(Qt.Horizontal)
        
        # 左侧面板 - 标注画布
        self.canvas = AnnotationCanvas()
        canvas_scroll = QScrollArea()
        canvas_scroll.setWidget(self.canvas)
        canvas_scroll.setWidgetResizable(True)
        
        # 设置滚动区域属性，以支持拖动
        canvas_scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAsNeeded)
        canvas_scroll.setVerticalScrollBarPolicy(Qt.ScrollBarAsNeeded)
        canvas_scroll.setFrameShape(QFrame.NoFrame)  # 无边框
        
        # 添加到拆分器
        splitter.addWidget(canvas_scroll)
        
        # 右侧面板 - 标注信息
        right_panel = QWidget()
        right_layout = QVBoxLayout(right_panel)
        
        # 标注类别选择
        class_group = QGroupBox("标注类别")
        class_layout = QVBoxLayout(class_group)
        
        class_select_layout = QHBoxLayout()
        self.class_combo = QComboBox()
        self.class_combo.setStyleSheet("QComboBox::drop-down { image: url(resources/arrow-down.svg); }")
        self.class_combo.addItem("person")
        self.class_combo.addItem("car")
        self.class_combo.addItem("dog")
        self.class_combo.currentTextChanged.connect(lambda text: self.change_class(text))
        self.class_combo.setContextMenuPolicy(Qt.CustomContextMenu)
        self.class_combo.customContextMenuRequested.connect(self.show_class_context_menu)
        class_select_layout.addWidget(self.class_combo)
        
        # 类别管理按钮布局
        class_manage_layout = QHBoxLayout()
        
        add_class_btn = QPushButton("添加类别")
        add_class_btn.clicked.connect(self.add_new_class)
        add_class_btn.setProperty("secondary", "true")
        class_manage_layout.addWidget(add_class_btn)
        
        edit_class_btn = QPushButton("编辑类别")
        edit_class_btn.clicked.connect(self.edit_class)
        edit_class_btn.setProperty("secondary", "true")
        class_manage_layout.addWidget(edit_class_btn)
        
        remove_class_btn = QPushButton("删除类别")
        remove_class_btn.clicked.connect(self.remove_class)
        remove_class_btn.setProperty("secondary", "true")
        class_manage_layout.addWidget(remove_class_btn)
        
        # 添加到类别组
        class_layout.addLayout(class_select_layout)
        class_layout.addLayout(class_manage_layout)
        right_layout.addWidget(class_group)
        
        # 图像信息
        info_group = QGroupBox("图像信息")
        info_layout = QVBoxLayout(info_group)
        
        self.image_info_label = QLabel("路径: ")
        self.image_size_label = QLabel("尺寸: ")
        self.annotation_count_label = QLabel("已标注: 0")
        
        info_layout.addWidget(self.image_info_label)
        info_layout.addWidget(self.image_size_label)
        info_layout.addWidget(self.annotation_count_label)
        right_layout.addWidget(info_group)
        
        # 标注列表
        annotations_group = QGroupBox("标注列表")
        annotations_layout = QVBoxLayout(annotations_group)
        
        self.annotations_list = QListWidget()
        self.annotations_list.currentItemChanged.connect(self.select_annotation)
        self.annotations_list.setContextMenuPolicy(Qt.CustomContextMenu)
        self.annotations_list.customContextMenuRequested.connect(self.show_annotation_context_menu)
        annotations_layout.addWidget(self.annotations_list)
        right_layout.addWidget(annotations_group)
        
        splitter.addWidget(right_panel)
        
        # 设置拆分器比例
        splitter.setSizes([int(self.width() * 0.7), int(self.width() * 0.3)])
        
        main_layout.addWidget(splitter)
        
        # 设置画布的类别选择器
        self.canvas.set_current_class(self.class_combo.currentText())
        
        # 连接画布的标注创建信号
        self.canvas.box_created.connect(self.add_annotation_to_list)
        
        # 加载项目的类别设置
        self.load_classes_from_project()
    
    def get_current_project_path(self):
        """获取当前项目路径，如果不在项目中则返回None"""
        try:
            # 获取主窗口
            main_window = self.window()
            if hasattr(main_window, 'current_project') and main_window.current_project:
                return main_window.current_project.get('path')
        except:
            pass
        return None
    
    def save_classes_to_project(self):
        """保存类别设置到project.json"""
        project_path = self.get_current_project_path()
        if not project_path:
            return
            
        try:
            import os
            import json
            
            # 获取项目配置文件路径
            config_path = os.path.join(project_path, "project.json")
            
            # 读取现有配置
            if os.path.exists(config_path):
                with open(config_path, 'r', encoding='utf-8') as f:
                    project_config = json.load(f)
            else:
                project_config = {}
            
            # 获取所有类别和颜色
            classes = []
            for i in range(self.class_combo.count()):
                class_name = self.class_combo.itemText(i)
                color = self.canvas.colors.get(class_name, Qt.red)
                classes.append({
                    "name": class_name,
                    "color": color.name()
                })
            
            # 更新配置
            project_config["classes"] = classes
            
            # 保存配置
            with open(config_path, 'w', encoding='utf-8') as f:
                json.dump(project_config, f, indent=4, ensure_ascii=False)
                
        except Exception as e:
            QMessageBox.warning(self, "警告", f"保存类别设置失败: {str(e)}")
    
    def load_classes_from_project(self):
        """从project.json加载类别设置"""
        project_path = self.get_current_project_path()
        if not project_path:
            return
            
        try:
            import os
            import json
            
            # 获取项目配置文件路径
            config_path = os.path.join(project_path, "project.json")
            
            # 读取配置
            if not os.path.exists(config_path):
                return
                
            with open(config_path, 'r', encoding='utf-8') as f:
                project_config = json.load(f)
            
            # 获取类别设置
            classes = project_config.get("classes", [])
            if not classes:
                return
                
            # 清空现有类别
            self.class_combo.clear()
            self.canvas.colors.clear()
            
            # 添加类别
            for class_info in classes:
                class_name = class_info["name"]
                color_name = class_info.get("color", "#ff0000")
                
                # 添加到下拉菜单
                self.class_combo.addItem(class_name)
                
                # 设置颜色
                self.canvas.colors[class_name] = QColor(color_name)
            
            # 如果没有加载到类别，添加默认类别
            if self.class_combo.count() == 0:
                self.class_combo.addItem("person")
                self.canvas.colors["person"] = QColor(255, 0, 0)
                self.class_combo.addItem("car")
                self.canvas.colors["car"] = QColor(0, 255, 0)
                self.class_combo.addItem("dog")
                self.canvas.colors["dog"] = QColor(0, 0, 255)
            
            # 选择第一个类别
            if self.class_combo.count() > 0:
                self.canvas.set_current_class(self.class_combo.itemText(0))
                
        except Exception as e:
            QMessageBox.warning(self, "警告", f"加载类别设置失败: {str(e)}")
            
            # 添加默认类别
            self.class_combo.clear()
            self.canvas.colors.clear()
            
            self.class_combo.addItem("person")
            self.canvas.colors["person"] = QColor(255, 0, 0)
            self.class_combo.addItem("car")
            self.canvas.colors["car"] = QColor(0, 255, 0)
            self.class_combo.addItem("dog")
            self.canvas.colors["dog"] = QColor(0, 0, 255)
            
            # 选择第一个类别
            self.canvas.set_current_class("person")
    
    def show_class_context_menu(self, pos):
        """显示类别下拉框的右键菜单"""
        menu = QMenu(self)
        
        edit_action = menu.addAction("编辑当前类别")
        edit_action.triggered.connect(self.edit_class)
        
        remove_action = menu.addAction("删除当前类别")
        remove_action.triggered.connect(self.remove_class)
        
        # 只有当有多于一个类别时才能删除
        if self.class_combo.count() <= 1:
            remove_action.setEnabled(False)
        
        menu.addSeparator()
        
        add_action = menu.addAction("添加新类别")
        add_action.triggered.connect(self.add_new_class)
        
        menu.exec_(self.class_combo.mapToGlobal(pos))
    
    def show_annotation_context_menu(self, pos):
        """显示标注列表的右键菜单"""
        if self.annotations_list.count() == 0:
            return
            
        selected_item = self.annotations_list.currentItem()
        if not selected_item:
            return
            
        menu = QMenu(self)
        
        remove_action = menu.addAction("删除此标注")
        index = self.annotations_list.row(selected_item)
        remove_action.triggered.connect(lambda: self.remove_annotation_at_index(index))
        
        edit_class_action = menu.addAction("修改类别")
        edit_class_action.triggered.connect(lambda: self.change_annotation_class(index))
        
        menu.addSeparator()
        
        remove_all_action = menu.addAction("删除所有标注")
        remove_all_action.triggered.connect(self.clear_annotations)
        
        menu.exec_(self.annotations_list.mapToGlobal(pos))
    
    def set_status(self, message, is_error=False, timeout=5000):
        """显示状态消息"""
        # 如果是错误消息，使用对话框显示
        if is_error:
            QMessageBox.warning(self, "警告", message)
        # 其他消息不显示，因为状态栏已被移除
        pass
        
    def open_image(self):
        """打开单个图像文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "打开图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp)"
        )
        
        if file_path:
            # 清除之前的标注
            if hasattr(self, 'canvas'):
                self.canvas.boxes = []
                self.annotations_list.clear()
            
            # 确保加载类别
            self.load_classes_from_project()
            
            # 检查当前项目，如果存在则复制图像到项目目录
            project_path = self.get_current_project_path()
            working_path = file_path
            
            if project_path:
                # 如果有项目，则复制图像到项目目录
                import shutil
                
                project_images_dir = os.path.join(project_path, "images")
                # 确保项目图像目录存在
                if not os.path.exists(project_images_dir):
                    os.makedirs(project_images_dir)
                
                # 复制图像到项目目录
                file_name = os.path.basename(file_path)
                target_path = os.path.join(project_images_dir, file_name)
                
                try:
                    # 复制文件（如果目标不存在或强制覆盖）
                    if not os.path.exists(target_path) or QMessageBox.question(
                        self, 
                        "文件已存在", 
                        f"文件 {file_name} 已存在于项目目录中，是否覆盖？",
                        QMessageBox.Yes | QMessageBox.No
                    ) == QMessageBox.Yes:
                        shutil.copy2(file_path, target_path)
                        working_path = target_path
                        self.set_status(f"已复制图像到项目目录: {file_name}")
                except Exception as e:
                    self.set_status(f"复制图像到项目失败: {str(e)}", True)
            
            # 清空当前图像文件列表，只包含此图像
            self.image_files = [working_path]
            self.current_index = 0
            self.current_dir = os.path.dirname(working_path)
            
            # 加载图像
            if self.canvas.set_image(working_path):
                # 更新UI
                self.image_info_label.setText(f"路径: {working_path}")
                self.image_size_label.setText(f"尺寸: {self.canvas.image.width()} x {self.canvas.image.height()}")
                self.annotation_count_label.setText(f"已标注: {len(self.canvas.boxes)}")
                self.annotations_list.clear()
                
                # 加载已有标注
                self.load_existing_annotations(working_path)
                
                # 禁用导航按钮，因为只有一张图片
                self.prev_action.setEnabled(False)
                self.next_action.setEnabled(False)
                
                # 设置焦点到画布，使快捷键能正常工作
                self.canvas.setFocus()
                
                self.set_status(f"已加载图像: {os.path.basename(working_path)}")
                
    def save_annotations(self):
        """保存当前标注结果"""
        if not hasattr(self.canvas, 'image_path'):
            self.set_status("没有打开的图像", True)
            return False
            
        # 获取图像路径的基本信息
        import os
        base_name = os.path.basename(self.canvas.image_path)
        base_name_without_ext = os.path.splitext(base_name)[0]
        
        # 确定标签目录
        image_dir = os.path.dirname(self.canvas.image_path)
        
        # 如果这是一个项目图像目录，则标签目录应该在同级
        if "images" in image_dir and image_dir.endswith("images"):
            # 项目结构，标签目录与图像目录同级
            project_dir = os.path.dirname(image_dir)
            labels_dir = os.path.join(project_dir, "labels")
        else:
            # 常规结构，标签目录在图像目录下
            labels_dir = os.path.join(image_dir, "labels")
        
        # 确保标签目录存在
        if not os.path.exists(labels_dir):
            os.makedirs(labels_dir)
            
        # 获取YOLO格式的标注
        annotations = self.canvas.export_annotations()
        
        # 保存到文件
        label_path = os.path.join(labels_dir, f"{base_name_without_ext}.txt")
        with open(label_path, 'w') as f:
            f.write('\n'.join(annotations))
            
        self.set_status(f"标注已保存: {base_name_without_ext}.txt")
        return True
        
    def clear_annotations(self):
        """清除所有标注"""
        self.canvas.clear_annotations()
        self.annotations_list.clear()
        self.annotation_count_label.setText("已标注: 0")
        
        # 清空后自动保存
        if hasattr(self.canvas, 'image_path'):
            self.save_annotations()
        
    def remove_annotation_at_index(self, index):
        """删除指定索引的标注"""
        if 0 <= index < len(self.canvas.boxes):
            # 从画布中删除
            self.canvas.boxes.pop(index)
            # 从列表中删除
            self.annotations_list.takeItem(index)
            # 更新计数
            self.annotation_count_label.setText(f"已标注: {len(self.canvas.boxes)}")
            # 更新画布显示
            self.canvas.update_pixmap()
        
    def change_annotation_class(self, index):
        """修改指定索引标注的类别"""
        if 0 <= index < len(self.canvas.boxes):
            rect, _, color = self.canvas.boxes[index]
            
            # 创建类别选择对话框
            dialog = QDialog(self)
            dialog.setWindowTitle("选择新类别")
            layout = QVBoxLayout(dialog)
            
            # 添加类别下拉框
            class_combo = QComboBox()
            for i in range(self.class_combo.count()):
                class_combo.addItem(self.class_combo.itemText(i))
            
            layout.addWidget(QLabel("选择新的类别:"))
            layout.addWidget(class_combo)
            
            # 添加按钮
            button_layout = QHBoxLayout()
            ok_button = QPushButton("确定")
            ok_button.clicked.connect(dialog.accept)
            cancel_button = QPushButton("取消")
            cancel_button.clicked.connect(dialog.reject)
            
            button_layout.addWidget(ok_button)
            button_layout.addWidget(cancel_button)
            
            layout.addLayout(button_layout)
            
            # 显示对话框
            if dialog.exec_() == QDialog.Accepted:
                new_class = class_combo.currentText()
                new_color = self.canvas.colors.get(new_class, Qt.red)
                
                # 更新标注
                self.canvas.boxes[index] = (rect, new_class, new_color)
                
                # 更新列表显示
                item_text = f"{new_class}: ({rect.x()}, {rect.y()}, {rect.width()}, {rect.height()})"
                self.annotations_list.item(index).setText(item_text)
                
                # 更新画布
                self.canvas.update_pixmap()
                
                # 保存自动保存标注
                if hasattr(self.canvas, 'image_path'):
                    self.save_annotations()
        
    def undo_annotation(self):
        """撤销最后一个标注"""
        if self.canvas.boxes:
            self.canvas.remove_last_annotation()
            if self.annotations_list.count() > 0:
                self.annotations_list.takeItem(self.annotations_list.count() - 1)
            self.annotation_count_label.setText(f"已标注: {len(self.canvas.boxes)}")
        
    def zoom_in_image(self):
        """放大图像"""
        # 获取鼠标当前位置
        cursor_pos = self.canvas.mapFromGlobal(QCursor.pos())
        # 如果鼠标在canvas范围内，则以鼠标位置为中心缩放
        if self.canvas.rect().contains(cursor_pos):
            self.canvas.zoom_in(cursor_pos)
        else:
            # 使用canvas中心点
            center_pos = QPoint(self.canvas.width() // 2, self.canvas.height() // 2)
            self.canvas.zoom_in(center_pos)
        
    def zoom_out_image(self):
        """缩小图像"""
        # 获取鼠标当前位置
        cursor_pos = self.canvas.mapFromGlobal(QCursor.pos())
        # 如果鼠标在canvas范围内，则以鼠标位置为中心缩放
        if self.canvas.rect().contains(cursor_pos):
            self.canvas.zoom_out(cursor_pos)
        else:
            # 使用canvas中心点
            center_pos = QPoint(self.canvas.width() // 2, self.canvas.height() // 2)
            self.canvas.zoom_out(center_pos)
    
    def reset_zoom(self):
        """重置图像缩放"""
        self.canvas.reset_zoom()
            
    def change_class(self, class_name):
        """更改当前标注类别"""
        self.canvas.set_current_class(class_name)
        
    def add_new_class(self):
        """添加新的标注类别"""
        class_name, ok = QInputDialog.getText(self, "新类别", "输入新类别名称:")
        if ok and class_name:
            # 检查是否已存在
            for i in range(self.class_combo.count()):
                if self.class_combo.itemText(i) == class_name:
                    QMessageBox.warning(self, "警告", f"类别 '{class_name}' 已存在！")
                    return
                    
            # 选择颜色
            color = QColorDialog.getColor(Qt.red, self, "选择类别颜色")
            if color.isValid():
                # 添加到下拉列表
                self.class_combo.addItem(class_name)
                self.class_combo.setCurrentText(class_name)
                
                # 设置颜色
                self.canvas.colors[class_name] = color
                
                # 保存类别到项目配置
                self.save_classes_to_project()
                
    def edit_class(self):
        """编辑当前选中的类别"""
        current_class = self.class_combo.currentText()
        if not current_class:
            return
            
        # 创建编辑对话框
        dialog = QDialog(self)
        dialog.setWindowTitle(f"编辑类别 '{current_class}'")
        layout = QVBoxLayout(dialog)
        
        # 添加类别名称输入
        name_layout = QHBoxLayout()
        name_layout.addWidget(QLabel("类别名称:"))
        name_edit = QLineEdit(current_class)
        name_layout.addWidget(name_edit)
        layout.addLayout(name_layout)
        
        # 添加颜色选择按钮
        current_color = self.canvas.colors.get(current_class, Qt.red)
        color_button = QPushButton("修改颜色")
        color_preview = QLabel()
        color_preview.setFixedSize(24, 24)
        color_preview.setStyleSheet(f"background-color: {current_color.name()};")
        
        # 创建一个可变的颜色变量来解决 nonlocal 问题
        selected_color = [current_color]  # 使用列表来允许内部函数修改
        
        def update_color():
            color = QColorDialog.getColor(selected_color[0], dialog, "选择类别颜色")
            if color.isValid():
                selected_color[0] = color  # 更新列表中的颜色
                color_preview.setStyleSheet(f"background-color: {color.name()};")
                
        color_button.clicked.connect(update_color)
        
        color_layout = QHBoxLayout()
        color_layout.addWidget(QLabel("类别颜色:"))
        color_layout.addWidget(color_preview)
        color_layout.addWidget(color_button)
        layout.addLayout(color_layout)
        
        # 添加按钮
        button_layout = QHBoxLayout()
        ok_button = QPushButton("确定")
        ok_button.clicked.connect(dialog.accept)
        cancel_button = QPushButton("取消")
        cancel_button.clicked.connect(dialog.reject)
        
        button_layout.addWidget(ok_button)
        button_layout.addWidget(cancel_button)
        
        layout.addLayout(button_layout)
        
        # 显示对话框
        if dialog.exec_() == QDialog.Accepted:
            new_class_name = name_edit.text().strip()
            
            # 检查输入是否有效
            if not new_class_name:
                QMessageBox.warning(self, "警告", "类别名称不能为空！")
                return
                
            # 检查是否与其他类别重名
            for i in range(self.class_combo.count()):
                if i != self.class_combo.currentIndex() and self.class_combo.itemText(i) == new_class_name:
                    QMessageBox.warning(self, "警告", f"类别 '{new_class_name}' 已存在！")
                    return
            
            # 获取选择的颜色
            new_color = selected_color[0]
            
            # 更新类别
            if new_class_name != current_class:
                # 更新已有标注中的类别名称
                for i in range(len(self.canvas.boxes)):
                    rect, class_name, color = self.canvas.boxes[i]
                    if class_name == current_class:
                        self.canvas.boxes[i] = (rect, new_class_name, new_color)
                        
                        # 更新列表项
                        item_text = f"{new_class_name}: ({rect.x()}, {rect.y()}, {rect.width()}, {rect.height()})"
                        self.annotations_list.item(i).setText(item_text)
                
                # 更新下拉框项
                current_index = self.class_combo.currentIndex()
                self.class_combo.setItemText(current_index, new_class_name)
                
                # 如果当前正在使用这个类别
                if self.canvas.current_class == current_class:
                    self.canvas.current_class = new_class_name
                
                # 更新颜色映射
                del self.canvas.colors[current_class]
                self.canvas.colors[new_class_name] = new_color
            else:
                # 仅更新颜色
                self.canvas.colors[current_class] = new_color
                
                # 更新已有标注的颜色
                for i in range(len(self.canvas.boxes)):
                    rect, class_name, _ = self.canvas.boxes[i]
                    if class_name == current_class:
                        self.canvas.boxes[i] = (rect, class_name, new_color)
            
            # 更新画布
            self.canvas.update_pixmap()
            
            # 保存类别到项目配置
            self.save_classes_to_project()
                
    def remove_class(self):
        """删除当前选中的类别"""
        # 检查是否只剩一个类别
        if self.class_combo.count() <= 1:
            QMessageBox.warning(self, "警告", "至少需要保留一个类别！")
            return
            
        current_class = self.class_combo.currentText()
        
        # 确认删除
        reply = QMessageBox.question(
            self, "确认删除",
            f"确定要删除类别 '{current_class}' 吗？\n\n注意：已标注为此类别的框将被删除！",
            QMessageBox.Yes | QMessageBox.No, QMessageBox.No
        )
        
        if reply == QMessageBox.Yes:
            # 删除相关标注
            new_boxes = []
            
            # 记录需要删除的索引
            indices_to_remove = []
            
            for i, (rect, class_name, color) in enumerate(self.canvas.boxes):
                if class_name != current_class:
                    new_boxes.append((rect, class_name, color))
                else:
                    indices_to_remove.append(i)
            
            # 更新标注框列表
            self.canvas.boxes = new_boxes
            
            # 从列表控件中移除，从后向前删除
            for i in sorted(indices_to_remove, reverse=True):
                self.annotations_list.takeItem(i)
            
            # 更新标注计数
            self.annotation_count_label.setText(f"已标注: {len(self.canvas.boxes)}")
            
            # 从下拉列表中删除
            current_index = self.class_combo.currentIndex()
            self.class_combo.removeItem(current_index)
            
            # 从颜色映射中删除
            if current_class in self.canvas.colors:
                del self.canvas.colors[current_class]
            
            # 更新画布
            self.canvas.update_pixmap()
            
            # 保存类别到项目配置
            self.save_classes_to_project()
    
    def add_annotation_to_list(self, rect, class_name):
        """将新标注添加到列表"""
        item_text = f"{class_name}: ({rect.x()}, {rect.y()}, {rect.width()}, {rect.height()})"
        self.annotations_list.addItem(item_text)
        self.annotation_count_label.setText(f"已标注: {len(self.canvas.boxes)}")
        
    def select_annotation(self, item):
        """选择列表中的标注项"""
        # 获取选中的索引
        index = self.annotations_list.row(item)
        if index >= 0 and index < len(self.canvas.boxes):
            # 高亮显示选中的框
            rect, class_name, color = self.canvas.boxes[index]
            
            # 暂时实现为更新画布，将选中的框绘制为更粗的边框
            # 在实际应用中可以添加更复杂的高亮效果
            self.canvas.update_pixmap()
        
    def export_dataset(self):
        """导出整个标注数据集"""
        # 检查是否有图像可以导出
        if not hasattr(self, 'image_files') or not self.image_files:
            if not hasattr(self.canvas, 'image_path'):
                self.set_status("请先加载图像！", True)
                return
                
        # 选择导出目录
        export_dir = QFileDialog.getExistingDirectory(self, "选择导出目录")
        if not export_dir:
            return
            
        # 创建数据集结构
        import os
        import shutil
        images_dir = os.path.join(export_dir, "images")
        labels_dir = os.path.join(export_dir, "labels")
        
        os.makedirs(images_dir, exist_ok=True)
        os.makedirs(labels_dir, exist_ok=True)
        
        # 保存所有图像和标注
        # 首先保存当前正在编辑的图像和标注（如果有）
        if hasattr(self.canvas, 'image_path'):
            # 保存当前标注
            self.save_annotations()
            
        # 创建进度对话框
        from PyQt5.QtWidgets import QProgressDialog
        from PyQt5.QtCore import Qt
        
        # 批量模式 - 导出所有已加载的图像和标注
        if hasattr(self, 'image_files') and self.image_files:
            total_files = len(self.image_files)
            
            # 显示进度对话框
            progress = QProgressDialog("正在导出数据集...", "取消", 0, total_files, self)
            progress.setWindowTitle("导出数据集")
            progress.setWindowModality(Qt.WindowModal)
            progress.setValue(0)
            
            # 计数器
            exported_images = 0
            exported_labels = 0
            
            # 导出所有图像和对应的标签
            for i, image_path in enumerate(self.image_files):
                # 检查是否取消
                if progress.wasCanceled():
                    break
                    
                # 更新进度
                progress.setValue(i)
                
                # 复制图像
                image_basename = os.path.basename(image_path)
                image_dest = os.path.join(images_dir, image_basename)
                try:
                    shutil.copy2(image_path, image_dest)
                    exported_images += 1
                except Exception as e:
                    continue
                    
                # 复制对应的标签文件
                base_name_without_ext = os.path.splitext(image_basename)[0]
                image_dir = os.path.dirname(image_path)
                
                # 确定标签文件路径
                if "images" in image_dir and image_dir.endswith("images"):
                    # 项目结构
                    project_dir = os.path.dirname(image_dir)
                    source_label_path = os.path.join(project_dir, "labels", f"{base_name_without_ext}.txt")
                else:
                    # 常规结构
                    source_label_path = os.path.join(image_dir, "labels", f"{base_name_without_ext}.txt")
                
                # 如果标签文件存在，则复制
                if os.path.exists(source_label_path):
                    label_dest = os.path.join(labels_dir, f"{base_name_without_ext}.txt")
                    try:
                        shutil.copy2(source_label_path, label_dest)
                        exported_labels += 1
                    except Exception as e:
                        pass
                
            # 完成进度
            progress.setValue(total_files)
            
            # 创建数据集配置文件
            dataset_yaml = os.path.join(export_dir, "dataset.yaml")
            class_names = [self.class_combo.itemText(i) for i in range(self.class_combo.count())]
            
            with open(dataset_yaml, 'w') as f:
                f.write(f"path: {export_dir}\n")
                f.write("train: images\n")
                f.write("val: images\n\n")
                f.write(f"nc: {len(class_names)}\n")
                f.write(f"names: {class_names}\n")
                
            self.set_status(f"数据集已导出: {exported_images}张图像, {exported_labels}个标签文件")
            
        else:
            # 单图模式 - 只导出当前图像
            if hasattr(self.canvas, 'image_path'):
                # 复制图像
                image_basename = os.path.basename(self.canvas.image_path)
                image_dest = os.path.join(images_dir, image_basename)
                shutil.copy2(self.canvas.image_path, image_dest)
                
                # 复制标签
                base_name_without_ext = os.path.splitext(image_basename)[0]
                annotations = self.canvas.export_annotations()
                label_path = os.path.join(labels_dir, f"{base_name_without_ext}.txt")
                with open(label_path, 'w') as f:
                    f.write('\n'.join(annotations))
                    
                # 创建数据集配置文件
                dataset_yaml = os.path.join(export_dir, "dataset.yaml")
                class_names = [self.class_combo.itemText(i) for i in range(self.class_combo.count())]
                
                with open(dataset_yaml, 'w') as f:
                    f.write(f"path: {export_dir}\n")
                    f.write("train: images\n")
                    f.write("val: images\n\n")
                    f.write(f"nc: {len(class_names)}\n")
                    f.write(f"names: {class_names}\n")
                    
                self.set_status(f"单张图像数据集已导出到: {export_dir}")
            else:
                self.set_status("没有数据可导出", True)
        
    def batch_annotate(self):
        """打开批量标注模式"""
        # 获取当前项目路径
        project_path = self.get_current_project_path()
        
        if project_path:
            # 如果有当前项目，默认复制图像到项目
            project_images_dir = os.path.join(project_path, "images")
            self.open_batch_mode(None, True, project_images_dir)
        else:
            # 没有当前项目，不复制
            self.open_batch_mode(None, False, None)
            
        # 设置焦点到画布，使快捷键能正常工作
        QTimer.singleShot(100, lambda: self.canvas.setFocus())
    
    def load_prev_image(self):
        """加载上一张图片"""
        if self.current_index > 0:
            self.load_image_at_index(self.current_index - 1)
    
    def load_next_image(self):
        """加载下一张图片"""
        if self.current_index < len(self.image_files) - 1:
            self.load_image_at_index(self.current_index + 1)
    
    def load_image_at_index(self, idx):
        """加载指定索引的图像"""
        if not self.image_files or idx < 0 or idx >= len(self.image_files):
            return False
            
        # 保存当前标注
        if hasattr(self.canvas, 'image_path'):
            self.save_annotations()
        
        # 确保加载类别
        self.load_classes_from_project()
        
        # 清除之前的标注
        self.canvas.boxes = []
        self.annotations_list.clear()
            
        # 加载新图像
        image_path = self.image_files[idx]
        if self.canvas.set_image(image_path):
            # 更新UI
            self.current_index = idx
            self.image_info_label.setText(f"图像 {idx+1}/{len(self.image_files)}")
            self.image_size_label.setText(f"尺寸: {self.canvas.image.width()} x {self.canvas.image.height()}")
            self.annotation_count_label.setText(f"已标注: {len(self.canvas.boxes)}")
            
            # 启用/禁用导航按钮
            self.prev_action.setEnabled(idx > 0)
            self.next_action.setEnabled(idx < len(self.image_files) - 1)
            
            # 加载已有标注
            self.load_existing_annotations(image_path)
            
            # 设置焦点到画布，使快捷键能正常工作
            self.canvas.setFocus()
            
            return True
        return False
    
    def load_existing_annotations(self, image_path):
        """加载已有的标注文件"""
        import os
        base_name = os.path.basename(image_path)
        base_name_without_ext = os.path.splitext(base_name)[0]
        image_dir = os.path.dirname(image_path)
        
        # 确定标签目录
        # 如果这是一个项目图像目录，则标签目录应该在同级
        if "images" in image_dir and image_dir.endswith("images"):
            # 项目结构，标签目录与图像目录同级
            project_dir = os.path.dirname(image_dir)
            label_path = os.path.join(project_dir, "labels", f"{base_name_without_ext}.txt")
        else:
            # 常规结构，标签目录在图像目录下
            label_path = os.path.join(image_dir, "labels", f"{base_name_without_ext}.txt")
        
        if not os.path.exists(label_path):
            return
            
        with open(label_path, 'r') as f:
            lines = f.readlines()
            
        img_width = self.canvas.image.width()
        img_height = self.canvas.image.height()
        
        # 类别索引映射
        class_indices = {}
        for i in range(self.class_combo.count()):
            class_name = self.class_combo.itemText(i)
            class_indices[i] = class_name
            
        # 处理每一行
        for line in lines:
            parts = line.strip().split()
            if len(parts) == 5:
                try:
                    class_idx = int(parts[0])
                    x_center = float(parts[1]) * img_width
                    y_center = float(parts[2]) * img_height
                    width = float(parts[3]) * img_width
                    height = float(parts[4]) * img_height
                    
                    # 转换为像素坐标
                    left = int(x_center - width / 2)
                    top = int(y_center - height / 2)
                    
                    # 创建矩形
                    rect = QRect(left, top, int(width), int(height))
                    
                    # 设置类别
                    class_name = class_indices.get(class_idx, "unknown")
                    
                    # 添加标注
                    color = self.canvas.colors.get(class_name, Qt.red)
                    self.canvas.boxes.append((rect, class_name, color))
                    self.add_annotation_to_list(rect, class_name)
                except:
                    continue
                    
        # 更新画布
        self.canvas.update_pixmap()
    
    def open_batch_mode(self, images_dir=None, copy_to_project=False, project_images_dir=None):
        """打开批量标注模式
        
        参数:
        images_dir - 图像目录路径，如果为None则弹出选择对话框
        copy_to_project - 是否复制图像到项目目录
        project_images_dir - 项目图像目录路径
        """
        # 如果未指定目录，弹出选择对话框
        if images_dir is None:
            images_dir = QFileDialog.getExistingDirectory(self, "选择图像目录")
            
        if not images_dir:
            return
            
        # 保存当前目录
        self.current_dir = images_dir
        
        # 清除之前的标注
        if hasattr(self, 'canvas'):
            self.canvas.boxes = []
            self.annotations_list.clear()
        
        # 确保从项目中加载类别定义
        self.load_classes_from_project()
        
        # 确定工作目录（原始目录或项目目录）
        working_dir = project_images_dir if copy_to_project and project_images_dir else images_dir
        
        # 创建标签目录（如果不存在）
        # 对于项目，labels目录应该与images目录在同一级
        parent_dir = os.path.dirname(working_dir)
        labels_dir = os.path.join(parent_dir, "labels")
        if not os.path.exists(labels_dir):
            os.makedirs(labels_dir)
        
        # 扫描目录中的所有图像
        self.image_files = []
        
        source_files = []
        for root, _, files in os.walk(images_dir):
            for file in files:
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
                    file_path = os.path.join(root, file)
                    source_files.append(file_path)
        
        # 如果需要复制图像到项目目录
        if copy_to_project and project_images_dir:
            # 确保项目图像目录存在
            if not os.path.exists(project_images_dir):
                os.makedirs(project_images_dir)
                
            # 创建进度对话框
            import shutil
            from PyQt5.QtCore import Qt
            
            # 计算需要处理的文件数量
            total_files = len(source_files)
            if total_files == 0:
                self.set_status("所选目录中没有找到支持的图像文件", True)
                return
                
            # 显示进度对话框
            progress = QProgressDialog("正在复制图像文件...", "取消", 0, total_files, self)
            progress.setWindowTitle("复制图像")
            progress.setWindowModality(Qt.WindowModal)
            progress.setValue(0)
            
            # 复制图像文件
            copied_count = 0
            for i, file_path in enumerate(source_files):
                # 检查用户是否取消
                if progress.wasCanceled():
                    break
                    
                # 更新进度
                progress.setValue(i)
                
                # 提取文件名并创建目标路径
                file_name = os.path.basename(file_path)
                target_path = os.path.join(project_images_dir, file_name)
                
                # 如果目标文件已存在，则跳过（避免重复复制）
                if os.path.exists(target_path):
                    continue
                    
                try:
                    # 复制文件
                    shutil.copy2(file_path, target_path)
                    self.image_files.append(target_path)
                    copied_count += 1
                except Exception as e:
                    pass
                    
            # 完成进度
            progress.setValue(total_files)
            
            # 更新状态
            self.set_status(f"已复制 {copied_count} 个图像文件到项目目录")
            
            # 如果没有复制任何文件，退出
            if not self.image_files:
                return
        else:
            # 不复制，直接使用原始路径
            self.image_files = source_files
        
        # 按文件名排序
        self.image_files.sort()
        
        # 如果找到图像，加载第一张
        if self.image_files:
            self.load_image_at_index(0)
            self.set_status(f"已加载{len(self.image_files)}张图像，使用上一张/下一张按钮浏览")
        else:
            self.set_status("所选目录中没有找到支持的图像文件", True)

    def keyPressEvent(self, event):
        """处理键盘事件"""
        # 阻止事件传播到父组件，确保快捷键在当前组件处理
        if event.key() == Qt.Key_Plus or event.key() == Qt.Key_Equal:
            # 放大
            self.zoom_in_image()
            event.accept()
            return
        elif event.key() == Qt.Key_Minus:
            # 缩小
            self.zoom_out_image()
            event.accept()
            return
        elif event.key() == Qt.Key_0:
            # 重置缩放
            self.reset_zoom()
            event.accept()
            return
        elif event.key() == Qt.Key_A:
            # 上一张图片
            if self.prev_action.isEnabled():
                self.load_prev_image()
                event.accept()
                return
        elif event.key() == Qt.Key_D:
            # 下一张图片
            if self.next_action.isEnabled():
                self.load_next_image()
                event.accept()
                return
        elif event.key() == Qt.Key_S and event.modifiers() & Qt.ControlModifier:
            # Ctrl+S 保存
            self.save_annotations()
            event.accept()
            return
        elif event.key() == Qt.Key_Z and event.modifiers() & Qt.ControlModifier:
            # Ctrl+Z 撤销
            self.undo_annotation()
            event.accept()
            return
        elif event.key() == Qt.Key_Delete and event.modifiers() & Qt.ControlModifier:
            # Ctrl+Delete 清除所有标注
            self.clear_annotations()
            event.accept()
            return
        elif event.key() == Qt.Key_O and event.modifiers() & Qt.ControlModifier:
            # Ctrl+O 打开图像
            self.open_image()
            event.accept()
            return
        elif event.key() == Qt.Key_B and event.modifiers() & Qt.ControlModifier:
            # Ctrl+B 批量标注
            self.batch_annotate()
            event.accept()
            return
        
        # 如果没有处理，则交给基类处理
        super().keyPressEvent(event)

    def load_project_images(self):
        """加载当前项目中的图片"""
        # 获取当前项目路径
        project_path = self.get_current_project_path()
        
        if not project_path:
            QMessageBox.warning(self, "警告", "没有打开的项目，请先打开或创建一个项目")
            return
            
        # 获取项目图片目录
        images_dir = os.path.join(project_path, "images")
        
        if not os.path.exists(images_dir):
            QMessageBox.warning(self, "警告", f"项目图片目录不存在: {images_dir}")
            return
            
        # 检查目录是否有图片
        image_files = []
        for ext in ['.jpg', '.jpeg', '.png', '.bmp']:
            image_files.extend(glob.glob(os.path.join(images_dir, f"*{ext}")))
            
        if not image_files:
            QMessageBox.warning(self, "警告", "项目图片目录中没有找到图片文件")
            return
            
        # 使用 open_batch_mode 加载图片，但不复制（因为已经在项目目录中）
        self.open_batch_mode(images_dir, copy_to_project=False)
        
        # 设置焦点到画布，使快捷键能正常工作
        QTimer.singleShot(100, lambda: self.canvas.setFocus())
        
        self.set_status(f"已加载项目图片: {len(self.image_files)}张")


# 测试代码
if __name__ == "__main__":
    import sys
    from PyQt5.QtWidgets import QApplication
    
    app = QApplication(sys.argv)
    window = AnnotationPanel()
    window.showMaximized()  # 设置为最大化显示
    sys.exit(app.exec_()) 